# 1. 导入依赖,定义类,传入数据
import numpy as np

class KMeansUtils:
    # 2. 初始化参数
    def __init__(self, data, cluster_num):
        self.data = data        # m*n
        self.cluster_num = cluster_num  # k

    # 3. 训练方法
    def train(self, max_iter=1000):
        # 3.1 初始化质心点
        centroids = self.init_centroids(self.data, self.cluster_num)    # k*2
        # 3.2 根据迭代次数循环执行
        for i in range(max_iter):
            # 3.3 计算所有点到k个质心的距离,找到最近最近的质心
            distances = self.compute_distances(self.data, centroids)    # m*k
            # 3.4 更新质心
            centroids = self.update_centroids(self.data,distances)    # k*2
        # 3.5 对举例进行处理，将举例最近的质心记录下来（m*k--->m*1）

        distances = distances.argmin(axis=1).reshape(-1, 1)    # m*1
        # 3.6 返回质心和相应的距离
        return centroids, distances

    # 4.初始化质心点函数
    def init_centroids(self,data,cluster_num):
        # 4.1 随机选择k个点
        centroids = data[np.random.choice(data.shape[0], cluster_num, replace=False)]   # k*2
        return centroids

    # 5. 计算到所有质心的距离函数
    def compute_distances(self,data, centroids):
        # 5.1 计算到所有质心的距离
        distances = np.zeros((data.shape[0], self.cluster_num))    # m*k
        # - 对每个样本进行遍历,判断到哪个质心最近,将质心存储到返回矩阵中
        for i in range(self.cluster_num):
            distances[:, i] = np.linalg.norm(self.data - centroids[i], axis=1)
        return distances

    # 6. 更新质心函数
    def update_centroids(self,data, distances):      # distances: m*k
        # 6.1 计算每个质心的新坐标
        centroids = np.zeros((self.cluster_num, data.shape[1]))    # k*2
        # 重新计算质心坐标,属于该质心范围内的坐标计算均值
        for i in range(self.cluster_num):
            centroids[i] = np.mean(data[distances[:, i].argmin()], axis=0)    # 找到距离最近的样本,计算均值作为新的质心坐标
        return centroids


